from queue import PriorityQueue
import numpy as np

def display(map, start, goal):
    vis = np.zeros_like(map, dtype=np.int32)
    vis[map] = 1
    vis[start] += 2
    vis[goal] += 4
    print(vis)

class MyKey:
    def __init__(self, pos, priority):
        self.pos = pos
        self.priority = priority

    def __lt__(self, other):
        return self.priority < other.priority


class MyPriorityQueue:
    def __init__(self):
        self.dict = {}
        self.pq = PriorityQueue()

    def peek(self):
        key = self.pop()
        if key != None:
            self.push(key)
        return key

    def pop(self):
        if self.pq.empty():
            return None
        key = self.pq.get()
        while not self.pq.empty() and (key.pos not in self.dict or key != self.dict[key.pos]):
            key = self.pq.get()
        if key != self.dict[key.pos]:
            return None
        del self.dict[key.pos]
        return key

    def push(self, key):
        self.dict[key.pos] = key
        self.pq.put(key)
        
    def update(self, key):
        self.push(key)

    def remove(self, pos):
        if pos in self.dict:
            del self.dict[pos]

    def contains(self, pos):
        return pos in self.dict

# This is based on the following paper.
# Koenig S, Likhachev M. D^* lite[J]. Aaai/iaai, 2002, 15: 476-483.
# https://www.aaai.org/Papers/AAAI/2002/AAAI02-072.pdf
class DStarLiteRouting:
    def __init__(self, map, start, goal):
        self.map = map
        self.infty = 10 * map.shape[0] * map.shape[1]
        self.pq = MyPriorityQueue()
        self.k_m = 0
        self.g = np.full_like(self.map, self.infty, dtype=np.int32)
        self.rhs = np.full_like(self.map, self.infty, dtype=np.int32)
        self.rhs[goal] = 0.
        self.pq.push(MyKey(goal, (self.h(start, goal), 0.)))
        self.start = start
        self.goal = goal
        self.last = start
        self.computeShortestPath()

    @staticmethod
    def h(u, v):
        ux, uy = u
        vx, vy = v
        dx = abs(ux - vx)
        dy = abs(uy - vy)
        # diagnoal
        # return max(dx, dy) + min(dx, dy) * 0.414
        # non-diagnoal
        return dx + dy

    def c(self, u, v):
        # assert v in self.neighbor(u)
        if self.map[u] or self.map[v]:
            return self.infty
        else:
            return 1
    
    def neighbor(self, x):
        # diagnoal
        # directions = [(-1, 0), (-1, 1), (0, 1), (1, 1), (1, 0), (1, -1), (0, -1), (-1, -1)]
        # non-diagnoal
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
        neighbors = []
        for direction in directions:
            newx = x[0] + direction[0]
            newy = x[1] + direction[1]
            if newx >= 0 and newx < self.map.shape[0] and newy >= 0 and newy < self.map.shape[1]:
                neighbors.append((newx, newy))
        return neighbors        

    def calKey(self, x):
        return MyKey(x, (min(self.g[x], self.rhs[x]) + self.h(self.start, x) + self.k_m, min(self.g[x], self.rhs[x])))

    def update(self, x):
        if self.g[x] != self.rhs[x] and self.pq.contains(x):
            self.pq.update(self.calKey(x))
        elif self.g[x] != self.rhs[x] and not self.pq.contains(x):
            self.pq.push(self.calKey(x))
        elif self.g[x] == self.rhs[x] and self.pq.contains(x):
            self.pq.remove(x)

    def computeShortestPath(self):
        while self.rhs[self.start] > self.g[self.start] or (self.pq.peek() != None and self.pq.peek() < self.calKey(self.start)):
            key = self.pq.pop()
            new_key = self.calKey(key.pos)
            if key < new_key:
                self.pq.update(new_key)
            elif self.g[key.pos] > self.rhs[key.pos]:
                self.g[key.pos] = self.rhs[key.pos]
                self.pq.remove(key.pos)
                for s in self.neighbor(key.pos):
                    if s != self.goal:
                        self.rhs[s] = min(self.rhs[s], self.c(s, key.pos) + self.g[key.pos])
                    self.update(s)
            else:
                g_old = self.g[key.pos]
                self.g[key.pos] = self.infty
                neighbors = self.neighbor(key.pos)
                neighbors.append(key.pos)
                for s in neighbors:
                    if self.rhs[s] == self.c(s, key.pos) + g_old and s != self.goal:
                        self.rhs[s] = self.infty
                        for ss in self.neighbor(s):
                            self.rhs[s] = min(self.rhs[s], self.c(s, ss) + self.g[ss])
                    self.update(s)

    def updateMap(self, points):
        self.k_m += self.h(self.last, self.start)
        self.last = self.start
        for u in points:
            self.map[u] = not self.map[u]
            neighbors = self.neighbor(u)
            neighbors.append(u)
            for v in neighbors:
                if v != self.goal:
                    self.rhs[v] = self.infty
                    for ss in self.neighbor(v):
                        self.rhs[v] = min(self.rhs[v], self.c(v, ss) + self.g[ss])
                self.update(v)
        self.computeShortestPath()
        display(self.map, self.start, self.goal)

    def assignTurning(self, direction, movement_callback):
        if self.start != self.goal:
            if self.rhs[self.start] >= self.infty:
                print('there is no known path')
                display(self.map, self.start, self.goal)
                return
            min_dis = self.infty
            min_next = (-1, -1)
            min_dir = -1
            for i, s in enumerate(self.neighbor(self.start)):
                dis = self.c(self.start, s) + self.g[s]
                if dis < min_dis:
                    min_next = s
                    min_dis = dis
                    min_dir = i
            if min_next == (-1, -1):
                print('there is no known path')
                display(self.map, self.start, self.goal)
                return
            if min_dir == direction:
                self.start = min_next
                movement_callback(0)
            else:
                movement_callback((min_dir - direction) % 4)
        else:
            print('The robot has arrived.')
            exit()
                